# -*- coding: utf-8 -*-
"""
#-------------------------------------------------------------------------------
# Nom:         gr4j
#
# Auteur:      Alain Gauthier
#
# Créé le:     22/03/2013
# Copyright:   (c) eaucéa 2014
# Licence:     GNU Public Licence (GPL)

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

Retranscription du modèle GR4J développé par IRSTEA:
c'est la fonction de transfert pluie-débit du modèle classique GR4J
http://fresno.cemagref.fr/webgr/Modeles/GR4J/fonctionnement_GR4J.htm

Ref:
Perrin, C., Michel, C., Andréassian, V.
Improvement of a parsimonious model for streamflow simulation
(2003) Journal of Hydrology, 279 (1-4), pp. 275-289.

#-------------------------------------------------------------------------------
"""

# modules externes
# ----------------
import numpy as np
import scipy.optimize as opt
import math

#-------------------------------------------------------------------------------

def gr4j(Ax,Apluie,Aevapo,AniveauIni,Apas=24):
    """
    Fonction de transfert pluie-débit définie selon le modèle GR4J de IRSTEA
    Paramètres de la fonction
    Apas = pas de temps (heures) : 24 valeur par défaut
    AniveauIni = vecteur contenant les niveaux de remplissage initiaux des réservoirs (en %)
    Apluie = série temporelle de pluie donnée au pas de temps Apas
    Aevapo = série temporelle d'ETP donnée au pas de temps Apas

    Paramètre du modèle GR4J :
    Ax[0] = capacité du réservoir de production (X1 > 1mm)
    Ax[1] = coefficient d'échanges souterrains  (X2 dans [-100..100] mm)
    Ax[2] = capacité du réservoir de routage    (X3 > 1mm)
    Ax[3] = temps de base del'hydrogramme unitaire (X4 en nombre de pas de temps (fixé pas Apas))

    Résultat : dictionnaire contenant les éléments suivants
    ['Q']   = série temporelle des débits calculés aux pas de temps Apas (en mm)
    ['S']   = série temporelle du volume de remplissage du réservoir de production (en mm)
    ['R']   = série temporelle du volume de remplissage du réservoir de routage (en mm)
    ['ECH'] = échanges d'eau
    ['EN']  = évaporation nette
    ['PN']  = pluie nette

    Ref:
    Perrin, C., Michel, C., André Assian, V.
    Improvement of a parsimonious model for streamflow simulation
    (2003) Journal of Hydrology, 279 (1-4), pp. 275-289.
    """

    # initialisation du résultat
    resultat = dict()
    resultat['Q'] = []
    resultat['S'] = []
    resultat['R'] = []
    resultat['ECH'] = []
    resultat['EN'] = []
    resultat['PN'] = []

    # --------------------------------------------------------------------
    # fonction principale de transfert

    # initialisations ----------------------------------------------------
    nVal = min(Apluie.size,Aevapo.size)
    Qcalc = -9.999*np.ones(nVal)
    Scalc = np.zeros(nVal)
    Rcalc = np.zeros(nVal)
    ECHcalc = np.zeros(nVal)
    ERcalc = np.zeros(nVal)
    PRcalc = np.zeros(nVal)

    # taille max des HU
    tailleHU = 1500
    if (Apas<=6)  : nH=int(tailleHU/3.)-1
    if (Apas>6)   : nH=int(tailleHU/25.)
    if (Apas>=24) : nH=int(tailleHU/50.)

    # percolation
    if (Apas>=12): PERCFACTOR=2.25 # Daily time steps
    if (Apas<12) : PERCFACTOR=4    # Hourly time steps

    # gestion des données manquantes
    NLACMAX = 5
    nLAC=0

    # paramètres de GR4J (copie locale) : X1,X2,X3 et X4 sont bornées par des valeurs limites
    Lx = Ax.copy()
    Lx[0] = max(Lx[0],1)
    Lx[1] = min(max(Lx[1],-100),100)
    Lx[2] = max(Lx[2],1)
    Lx[3] = min(max(Lx[3],0.5),nH)

    # remplissage initial des réservoirs
    LniveauIni = AniveauIni.copy()
    if(LniveauIni[0]<=0): LniveauIni[0]=0.001
    if(LniveauIni[0]>1) : LniveauIni[0]=1
    if(LniveauIni[1]<=0): LniveauIni[1]=0.001
    if(LniveauIni[1]>1) : LniveauIni[1]=1

    # niveau de remplissage
    Sini=Lx[0]*LniveauIni[0];
    Rini=Lx[2]*LniveauIni[1];

    # calcul des hydrogrammes unitaires
    LHU1 = HU1 (nH,Lx[3],Apas)
    LHU2 = HU2 (nH,Lx[3],Apas)
    Q9 = np.zeros(nH)
    Q1 = np.zeros(2*nH)

    # boucle sur le pas de temps
    S = Sini
    R = Rini
    for itps in range (nVal):
        P=Apluie[itps]; E=Aevapo[itps]
        # gestion des données manquantes
        if((P<0)&(nLAC<=NLACMAX)):
            P=0; nLAC = nLAC +1
        if((E<0)&(nLAC<=NLACMAX)):
            E=0; nLAC = nLAC +1
        if(nLAC>NLACMAX):
            P=0; E=0; nLAC=0

        # fonction de transfert au pas de temps itps

        # ......Production.............
        if(P>E):
            EN=0.0
            ES=0.0
            WS=(P-E)/Lx[0]
            if(WS>=13):WS=13.0
            PS=Lx[0]*(1-(S/Lx[0])**2)*math.tanh(WS)/(1+S/Lx[0]*math.tanh(WS))
            PR=P-E-PS
        else:
            WS=(E-P)/Lx[0]
            if(WS>=13):WS=13.0
            ES=S*(2-S/Lx[0])*math.tanh(WS)/(1+(1-S/Lx[0])*math.tanh(WS))
            PS=0.0
            PR=0.0
            EN=E-P
        # mise à jour du niveau du réservoir de production
        S+=PS-ES

        # ......Percolation.............
        PERC=S/pow(1+(S/(PERCFACTOR*Lx[0]))**4,0.25)
        PERC=S-PERC
        S=S-PERC

        PR+=PERC

        # ......UH1 ....................
        for k in range(nH-1):
            Q9[k]=Q9[1+k]+0.9*LHU1[k]*PR
        Q9[nH-1]=0.9*LHU1[nH-1]*PR

        # ......UH2 ....................
        for l in range(2*nH-1):
            Q1[l]=Q1[1+l]+0.1*LHU2[l]*PR
        Q1[2*nH-1]=0.1*LHU2[2*nH-1]*PR

        # ......Echange souterrain ....................
        ECH=Lx[1]*(R/Lx[2])**3.5

        # ......niveau du réservoir de routage ..........
        TP=R+Q9[0]+ECH

        # Si le réservoir ne contient pas assez
        ech1=-R-Q9[0];R=0

        if(TP>=0): R=TP;ech1=ECH

        QR=R/pow(1+(R/Lx[2])**4,0.25)
        QR=R-QR
        R=R-QR

        # ......ruissellement direct........
        QD=0

        # si HU ne fournit pas assez d'eau
        TP=Q1[0]+ECH
        ech2 = -Q1[0]
        if(TP>0): QD=TP;ech2=ECH

        # ......Ecoulement total..................
        Q=QD+QR;

        #--- résultats
        Qcalc[itps]=Q
        Scalc[itps]=S
        Rcalc[itps]=R
        ECHcalc[itps]=ech1+ech2
        ERcalc[itps]=ES+EN
        PRcalc[itps]=PR

        #--- gestion des données manquantes (valeurs initiales S et R)----------
        if(nLAC<=NLACMAX):
            Sini=S
            Rini=R
    # fin de boucle sur les pas de temps
    resultat['Q'] = Qcalc
    resultat['S'] = Scalc
    resultat['R'] = Rcalc
    resultat['ECH'] = ECHcalc
    resultat['EN'] = ERcalc
    resultat['PN'] = PRcalc

    return resultat

# --------------------------------------------------------------------

def gr4j_initialisation(AX3,AQj0):
    """
    Initialisation pour vérifier que Qr=AQj0 la condition initiale de la chronique de débit à satisfaire
    Paramètres de la fonction:
    Ax   = Paramètre du modèle GR4J (détail dans la fonction gr4j)
    AQj0 = Débit recheché

    Resultat : remplissage du réservoir de routage (en mm) permettant d'obtenir un ruissellement AQj0
    """

    # paramètres de GR4J (copie locale) : X1,X2,X3 et X4 sont bornées par des valeurs limites
    LX3 = AX3
    LX3 = max(LX3,1)

    # le résultat est le zéro de la fonction fQrQ0 permettant d'ajuster la valeur initiale de débit Qr=AQj0
    Rini = LX3 # condition initiale
    resultat = opt.newton(fQrQ0,Rini,dfQrQ0,(LX3,AQj0))
    return resultat

# --------------------------------------------------------------------

def fQr(AR,AX3):
    """
    fonction de vidange du réservoir de routage de GR4J
    AR : volume de remplissage du réservoir
    AX3 : capacité du réservoir
    """
    y= (AR/AX3)**4
    return AR*(1-1/pow(1+y,0.25))

def fQrQ0(AR,AX3,AQ0):
    """
    fonction renvoyant fQr(R)-Q0 dont il faut trouver la racine R0
    """
    return fQr(AR,AX3)-AQ0

def dfQrQ0(AR,AX3,AQ0):
    """
    fonction dérivée de fQrQ0 par rapport à AR
    """
    y= (AR/AX3)**4
    return 1-1/pow(1+y,0.25)+y/pow(1+y,1.25)


# --------------------------------------------------------------------
# définition des hydrogrammes unitaires HU1 et HU2
def SS1(Aindice,ANMax,Apas):
    """
    fonction en S pour calcul de HU1
    paramètres
    Aindice : indice de pas de temps de calcul
    ANMax   : indice de pas de temps max
    Apas : valeur du pas de temps (en heures)
    """

    # Exposant de l'UH
    Expos=2.5;
    if (Apas<6): Expos=1.25

    if (Aindice<0):
        return 0
    if (Aindice < ANMax):
        return (Aindice/ANMax)**Expos
    if (Aindice >=ANMax):
        return 1

# --------------------------------------------------------------------

def SS2(Aindice,ANMax,Apas):
    """
    fonction en S pour calcul de HU2
    paramètres
    Aindice : indice de pas de temps de calcul
    ANMax   : indice de pas de temps max
    Apas : valeur du pas de temps (en heures)
    """

    # Exposant de l'UH
    Expos=2.5;
    if (Apas<6): Expos=1.25

    if (Aindice<0):
        return 0
    if (Aindice <= ANMax):
        return 0.5*(Aindice/ANMax)**Expos
    if (Aindice < 2*ANMax):
        return 1-0.5*(2-Aindice/ANMax)**Expos
    if (Aindice >=2*ANMax):
        return 1

# --------------------------------------------------------------------

def HU1 (ANH,ANMax,Apas):
    """ hydrogramme unitaire HU1 """
    resultat = np.zeros(ANH)
    for i in range(ANH):
        resultat[i] = SS1(i+1,ANMax,Apas)-SS1(i,ANMax,Apas)
    return resultat

def HU2 (ANH,ANMax,Apas):
    """ hydrogramme unitaire HU2 """
    resultat = np.zeros(2*ANH)
    for i in range(2*ANH):
        resultat[i] = SS2(i+1,ANMax,Apas)-SS2(i,ANMax,Apas)
    return resultat

# --------------------------------------------------------------------
# --------------------------------------------------------------------
